package com.qzdatasoft;

import com.qzdatasoft.ratelimit.bean.ResourceKey;
import com.qzdatasoft.ratelimit.adapter.RateLimitRuleCacheManager;
import com.qzdatasoft.ratelimit.core.LeakyBucketLimiter;
import com.qzdatasoft.ratelimit.core.CustomRateLimitRule;
import com.qzdatasoft.ratelimit.core.RequestTask;
import com.qzdatasoft.ratelimit.bean.RateLimitRule;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.concurrent.ConcurrentHashMap;

/**
 * 分层限流策略管理器
 * 模仿Nginx的limit_req_zone分层配置
 */
public class HierarchicalRateLimiter {
    private static final Logger logger = LoggerFactory.getLogger(RateLimitRuleCacheManager.class);

    // 配置管理系统
//    private final HierarchicalRateLimiterConfig config;

    // 第一层：全局限流器（系统保护）
    private final LeakyBucketLimiter globalBucket;

    // 第二层：业务级限流器（API/服务隔离）
    // key = "business:+业务关键字如ip", value = TimerDrivenLeakyBucket
    private final ConcurrentHashMap<String, LeakyBucketLimiter> businessBuckets;

    // 第三层：用户级限流器（公平性保障）
    // key = "user:+用户关键字如userId", value = TimerDrivenLeakyBucket
    private final ConcurrentHashMap<String, LeakyBucketLimiter> userBuckets;

    // 清理过期桶的定时器
    //private final ConcurrentHashMap<String, AtomicLong> bucketLastAccess;
    @Autowired
    private RateLimitRuleCacheManager cacheManager;
    /**
     * 创建分层限流器
     */
    public HierarchicalRateLimiter() {
        //this.config = new HierarchicalRateLimiterConfig();
        this.businessBuckets = new ConcurrentHashMap<>();
        this.userBuckets = new ConcurrentHashMap<>();
        //this.bucketLastAccess = new ConcurrentHashMap<>();

        // 初始化全局限流器（暂未使用）
        this.globalBucket = null;//new LeakyBucketLimiter(globalRate, globalBurst);

        logger.info("{} initialized.", this.getClass().getSimpleName());
    }

    // 返回 true 表示被限流
    public boolean tryAcquire(RequestTask task, String url, String ip, String user){
        boolean result =false;

        // key 采用固定的命名规范，按顺序依次组成各部分按优先级为： urlPattern[:ip[:user]]
        // 本版本只支持user+user，因此格式为 url[:user]

        // 检查规则， 按优先级逐级检查，数字越大，优先级越大，优先级定义：
        // 4. url + ip + user  -- 暂不支持
        // 3. url + ip         -- 暂不支持
        // 2. url + user
        // 1. url

        // priority level 4 - 基于 url + + ip + user 的限流检查

        // priority level 3 - 基于 url + ip 的限流检查

        // priority level 2 - 基于 url + user 的限流检查
        String resourceKeyStr="";
        RateLimitRule rule=null;
        try(ResourceKey resourceKey = getResourceKey(2, url, ip, user)) {
            resourceKeyStr = resourceKey.buildKey();
            rule = cacheManager.getRule(resourceKeyStr);
            if( null == rule){
                try(ResourceKey resourceKey1 = getResourceKey(2, url, ip, "*")) {
                    resourceKeyStr = resourceKey1.buildKey();
                    rule = cacheManager.getRule(resourceKeyStr);
                }
            }
        } catch (Exception e) {
            logger.error("HierarchicalRateLimiter tryAcquire on level 1 error for url: {}, ip: {}, user: {}", url, ip, user, e);
        }

        if( null == rule){
            return true; // 没有定义规则，不限流
        }

        CheckRuleType checkRuleType = tryAcquire(task, rule);
        if (!checkRuleType.equals(CheckRuleType.CONTINUE)) {
            return checkRuleType.equals(CheckRuleType.FALSE) ? false : true;
        }

        // priority level 1 - 基于 user 的全局限流检查
        try(ResourceKey resourceKey = getResourceKey(1, url, ip, user)) {
            return tryAcquire(task, rule).equals(CheckRuleType.FALSE) ? false : true;
        } catch (Exception e) {
            logger.error("HierarchicalRateLimiter tryAcquire on level 0 error for url: {}, ip: {}, user: {}", url, ip, user, e);
            return true; // 出错则默认放行
        }
    }

//    private CheckRuleType tryAcquire1(RequestTask task, String url, String ip, String user){
//        try(ResourceKey resourceKey = getResourceKey(2, url, ip, user)) {
//            CheckRuleType checkRuleType = tryAcquire(task, resourceKey.buildKey());
//            if (!checkRuleType.equals(CheckRuleType.CONTINUE)) {
//                return checkRuleType;
//            }
//        } catch (Exception e) {
//            logger.error("HierarchicalRateLimiter tryAcquire on level 1 error for url: {}, ip: {}, user: {}", url, ip, user, e);
//        }
//    }

    /**
     * 获取统计信息
     */
    public String getStats() {
        return String.format(
            "HierarchicalRateLimiter{全局桶=%s, 业务桶数=%d, 用户桶数=%d}",
            null==globalBucket?null:globalBucket.toString(),
            null==businessBuckets?null:businessBuckets.size(),
            userBuckets.size()
        );
    }

    /**
     * 关闭限流器，释放资源
     */
    public void shutdown() {
        globalBucket.shutdown();
        businessBuckets.values().forEach(LeakyBucketLimiter::shutdown);
        userBuckets.values().forEach(LeakyBucketLimiter::shutdown);
    }

    // 检查限流结果
    private CheckRuleType tryAcquire(RequestTask task, RateLimitRule rule){
        //RateLimitRuleDto rule = cacheManager.getRule(key);
        if( null == rule){
            return CheckRuleType.CONTINUE;// 未获取到, 无需限流
        }

        // 获取到了
        LeakyBucketLimiter bucket = userBuckets.computeIfAbsent(rule.getResourceKey(), k -> {
            CustomRateLimitRule customProps = new CustomRateLimitRule();
            //customProps.setBurstCapacity(rule.getBurstCapacity());

            long intervalNanos = (long) (1_000_000_000L / (rule.getRatePerSecond()*(1 + rule.getCompensationRate()))); // 地税时间间隔，注意单位纳秒
            customProps.setIntervalNanos(intervalNanos);

            int capacity = (int)intervalNanos/1000000; //容量默认设置为速率，注意速率单位为毫秒
            if(capacity<1)capacity=1;
            customProps.setCapacity(capacity+rule.getBurstCapacity());

            long timeoutMillis = rule.getTimeoutMillis()*1000000L; // 默认超时时间100毫秒
            customProps.setTimeoutNanos(timeoutMillis); // 转为纳秒

            return new LeakyBucketLimiter(customProps);
        });
        task.setLimiter(bucket);

        return userBuckets.get(rule.getResourceKey()).tryEnqueue(task, rule.getResourceKey())
                ? CheckRuleType.TRUE
                : CheckRuleType.FALSE;

    }

    // 生成不同优先级的 resource key
    private static final ResourceKey getResourceKey(int priority, String url, String ip, String user){
        return ResourceKey.builder()
                .priority(priority)
                .dimension(priority ==4 ? "ip_user" :
                           priority ==3 ? "ip" :
                           priority ==2 ? "user" :
                           "global")
                .urlPattern(url)
                .ip(ip)
                .user(user)
                .build();
    }

    // 检查规则结果类型
    private enum CheckRuleType{
        CONTINUE,  // 继续检查下一级规则
        FALSE, // 被限流
        TRUE  // 通过限流
    }

}